import flair
from flair.datasets import CONLL_03
from flair.embeddings import TransformerWordEmbeddings, TransformerDocumentEmbeddings
from flair.models import BinderModel
from flair.trainers import ModelTrainer

flair.set_seed(42)

corpus = CONLL_03().downsample(0.2)
print(corpus)

label_type = 'ner'
label_dict = corpus.make_label_dictionary(label_type=label_type, add_unk=False)
print(label_dict)

token_embeddings = TransformerWordEmbeddings(
    model='bert-base-uncased',
    fine_tune=True,
    subtoken_pooling="first_last",
    is_document_embedding=True
)

label_embeddings = TransformerDocumentEmbeddings(model='bert-base-uncased', fine_tune=True)

model = BinderModel(
    token_encoder=token_embeddings,
    label_encoder=label_embeddings,
    label_dictionary=label_dict,
    label_type=label_type
)

# 6. initialize trainer
trainer = ModelTrainer(model, corpus)

# 7. run fine-tuning
trainer.fine_tune('resources/taggers/binder-conll',
                  learning_rate=5.0e-5,
                  mini_batch_size=4,
                  mini_batch_chunk_size=4,  # remove this parameter to speed up computation if you have a big GPU
                  )